#include <cstdlib>
#include <cstring>
#include <cmath>
#include <ctime>
#include <random>
#include <iostream>
#include <algorithm>
#include "tensor_mat19.h"

namespace TensorMat19 {

// GridParams implementation
GridParams::GridParams(int res, bool is_global) : resolution(res), is_global_grid(is_global) {
    // Calculate grid dimensions based on resolution
    imt = 360 / resolution;
    jmt = 180 / resolution + 1;
    kmt = 88;
    
    // Calculate total number of grid points
    total_inner_points = imt * jmt * kmt;
    total_padded_points = (imt + 2) * (jmt + 2) * (kmt == 1 ? 1 : (kmt + 2));
}

int GridParams::inner_index(int i, int j, int k) const {
    // Convert to 0-based indexing for the inner grid (no padding)
    return i + j * imt + k * imt * jmt;
}

int GridParams::padded_index(int i, int j, int k) const {
    // Add 1 to account for padding
    i += 1;
    j += 1;
    
    // Calculate padded dimensions
    int padded_imt = imt + 2;
    int padded_jmt = jmt + 2;
    
    // For kmt=1 (2D case), don't add padding in z-direction
    if (kmt == 1) {
        return i + j * padded_imt;
    }
    
    // Otherwise, add padding in all dimensions
    k += 1;
    int padded_kmt = kmt + 2;
    
    // Calculate linear index
    return i + j * padded_imt + k * padded_imt * padded_jmt;
}

void GridParams::get_coords(int linear_idx, int& i, int& j, int& k) const {
    // Convert linear index to grid coordinates (inner grid)
    k = linear_idx / (imt * jmt);
    int remainder = linear_idx % (imt * jmt);
    j = remainder / imt;
    i = remainder % imt;
}

void GridParams::print_info() const {
    std::cout << "Grid parameters:" << std::endl;
    std::cout << "  Resolution: " << resolution << std::endl;
    std::cout << "  Grid type: " << (is_global_grid ? "Global (lat-lon)" : "Local (Cartesian)") << std::endl;
    std::cout << "  Dimensions: imt=" << imt << ", jmt=" << jmt << ", kmt=" << kmt << std::endl;
    std::cout << "  Inner points: " << total_inner_points << std::endl;
    std::cout << "  Padded points: " << total_padded_points << std::endl;
}

// TensorMatrix implementation
TensorMatrix::TensorMatrix(const GridParams& grid_params)
    : num_diags(19), grid(grid_params) {
    
    // Allocate memory for the tensor
    size = num_diags * grid.imt * grid.jmt * grid.kmt;
    data = new float[size]();  // Initialize to zeros
}

TensorMatrix::~TensorMatrix() {
    delete[] data;
}

float& TensorMatrix::at(int m, int i, int j, int k) {
    // Access element (m,i,j,k) in the tensor
    // m is 1-indexed according to task description (1-19)
    // i,j,k are 0-indexed
    
    // Convert m to 0-indexed
    int m_idx = m - 1;
    
    // Calculate linear index in the tensor
    int idx = m_idx + i * num_diags + j * num_diags * grid.imt + k * num_diags * grid.imt * grid.jmt;
    
    return data[idx];
}

float TensorMatrix::at(int m, int i, int j, int k) const {
    // Const version
    int m_idx = m - 1;
    int idx = m_idx + i * num_diags + j * num_diags * grid.imt + k * num_diags * grid.imt * grid.jmt;
    return data[idx];
}

void TensorMatrix::init_random() {
    // Initialize with random values using modern C++ random facilities
    // std::random_device rd;
    // std::mt19937 gen(rd());
    std::mt19937 gen(42);
    std::uniform_real_distribution<float> dis(0.0, 1.0);
    
    // Ensure diagonal dominance for better convergence
    for (int k = 0; k < grid.kmt; k++) {
        for (int j = 0; j < grid.jmt; j++) {
            for (int i = 0; i < grid.imt; i++) {
                // Generate random values for all 19 diagonals
                float diagonal_sum = 0.0f;
                
                // First fill non-diagonal elements (2-19)
                for (int m = 2; m <= 19; m++) {
                    float val = dis(gen);
                    at(m, i, j, k) = val;
                    diagonal_sum += std::abs(val);
                }
                
                // Make the diagonal dominant (positive value larger than sum of others)
                at(1, i, j, k) = diagonal_sum * 1.2f;
            }
        }
    }
}

void TensorMatrix::multiply(const float* x, float* y) const {
    // Matrix-vector multiplication for TensorMat19: y = A*x
    // x and y are padded vectors with shape (imt+2, jmt+2, kmt+2)
    
    // Define the 19-point stencil offsets
    // Using the order described in task: center, x-dir, y-dir, corners, z-dir
    // These offsets correspond to the indices in the tensor format
    const int stencil_offsets[19][3] = {
        {0, 0, 0},   // 1: center
        {-1, 0, 0},  // 2: x-direction (left)
        {1, 0, 0},   // 3: x-direction (right)
        {0, -1, 0},  // 4: y-direction (front)
        {0, 1, 0},   // 5: y-direction (back)
        {1, 1, 0},   // 6: corner (right-back)
        {1, -1, 0},  // 7: corner (right-front)
        {-1, -1, 0}, // 8: corner (left-front)
        {-1, 1, 0},  // 9: corner (left-back)
        {0, 0, -1},  // 10: z-direction (below)
        {-1, 0, -1}, // 11: edge in below layer (left)
        {1, 0, -1},  // 12: edge in below layer (right)
        {0, -1, -1}, // 13: edge in below layer (front)
        {0, 1, -1},  // 14: edge in below layer (back)
        {0, 0, 1},   // 15: z-direction (above)
        {-1, 0, 1},  // 16: edge in above layer (left)
        {1, 0, 1},   // 17: edge in above layer (right)
        {0, -1, 1},  // 18: edge in above layer (front)
        {0, 1, 1}    // 19: edge in above layer (back)
    };
    
    // Initialize result vector y to zero
    std::memset(y, 0, grid.total_padded_points * sizeof(float));
    
    // Perform the multiplication for all inner grid points
    for (int k = 0; k < grid.kmt; k++) {
        for (int j = 0; j < grid.jmt; j++) {
            for (int i = 0; i < grid.imt; i++) {
                float sum = 0.0f;
                
                // Apply the 19-point stencil
                for (int m = 1; m <= 19; m++) {
                    // Calculate index into padded vector x
                    int ni = i + stencil_offsets[m-1][0];
                    int nj = j + stencil_offsets[m-1][1];
                    int nk = k + stencil_offsets[m-1][2];
                    
                    // Get matrix coefficient from tensor format
                    float coef = at(m, i, j, k);
                    
                    // Get corresponding x value from padded vector
                    // Adding 1 to convert from inner grid to padded grid coordinates
                    int x_idx = grid.padded_index(ni, nj, nk);
                    
                    sum += coef * x[x_idx];
                }
                
                // Store result in padded vector y
                int y_idx = grid.padded_index(i, j, k);
                y[y_idx] = sum;
            }
        }
    }
}

void TensorMatrix::to_csr(CsrMatrix& csr) const {
    // Convert the tensor matrix to CSR format
    
    // Calculate number of rows and maximum non-zero elements
    int rows = grid.total_inner_points;
    int max_nnz = rows * 19;  // Each row can have up to 19 non-zero elements
    
    // Allocate memory for CSR format
    csr.rows = rows;
    csr.row_off = new int[rows + 1];
    csr.cols = new int[max_nnz];
    csr.data = new double[max_nnz];
    
    // Initialize row offsets
    csr.row_off[0] = 0;
    int nnz_count = 0;
    
    // Stencil offsets (same as in multiply method)
    const int stencil_offsets[19][3] = {
        {0, 0, 0},   {-1, 0, 0},  {1, 0, 0},   {0, -1, 0},  {0, 1, 0},
        {1, 1, 0},   {1, -1, 0},  {-1, -1, 0}, {-1, 1, 0},  {0, 0, -1},
        {-1, 0, -1}, {1, 0, -1},  {0, -1, -1}, {0, 1, -1},  {0, 0, 1},
        {-1, 0, 1},  {1, 0, 1},   {0, -1, 1},  {0, 1, 1}
    };
    
    // Process all inner grid points to build CSR matrix
    for (int k = 0; k < grid.kmt; k++) {
        for (int j = 0; j < grid.jmt; j++) {
            for (int i = 0; i < grid.imt; i++) {
                // Current row index
                int row = grid.inner_index(i, j, k);
                
                // For each point in the stencil
                for (int m = 1; m <= 19; m++) {
                    // Calculate neighbor coordinates
                    int ni = i + stencil_offsets[m-1][0];
                    int nj = j + stencil_offsets[m-1][1];
                    int nk = k + stencil_offsets[m-1][2];
                    
                    // Skip if out of bounds (boundary cells)
                    if (ni < 0 || ni >= grid.imt || nj < 0 || nj >= grid.jmt || nk < 0 || nk >= grid.kmt) {
                        continue;
                    }
                    
                    // Calculate column index
                    int col = grid.inner_index(ni, nj, nk);
                    
                    // Add to CSR format
                    csr.cols[nnz_count] = col;
                    csr.data[nnz_count] = static_cast<double>(at(m, i, j, k));
                    nnz_count++;
                }
                
                // Set row offset for next row
                csr.row_off[row + 1] = nnz_count;
            }
        }
    }
    
    // Set actual data size
    csr.data_size = nnz_count;
}

// Boundary communication for vectors
void communicate_boundaries(float* vec, const GridParams& grid) {
    // Apply boundary conditions to the padded vector
    
    if (grid.is_global_grid) {
        // Global grid: Periodic in x-direction, Mirror in y-direction
        
        // Periodic boundary in x-direction
        for (int k = 0; k < grid.kmt; k++) {
            for (int j = 0; j < grid.jmt; j++) {
                // Left boundary <- Right inner
                vec[grid.padded_index(-1, j, k)] = vec[grid.padded_index(grid.imt - 1, j, k)];
                
                // Right boundary <- Left inner
                vec[grid.padded_index(grid.imt, j, k)] = vec[grid.padded_index(0, j, k)];
            }
        }
        
        // Mirror boundary in y-direction
        for (int k = 0; k < grid.kmt; k++) {
            for (int i = 0; i < grid.imt; i++) {
                // Front boundary <- Front inner reflected
                vec[grid.padded_index(i, -1, k)] = vec[grid.padded_index(i, 0, k)];
                
                // Back boundary <- Back inner reflected
                vec[grid.padded_index(i, grid.jmt, k)] = vec[grid.padded_index(i, grid.jmt - 1, k)];
            }
        }
    } else {
        // Local grid: Dirichlet boundary condition (zeros at boundaries)
        // No need to communicate boundaries, they remain zero
    }
    
    // For both grid types, set z-direction boundaries (only if kmt > 1)
    if (grid.kmt > 1) {
        for (int j = 0; j < grid.jmt; j++) {
            for (int i = 0; i < grid.imt; i++) {
                // Bottom boundary
                vec[grid.padded_index(i, j, -1)] = 0.0f;
                
                // Top boundary
                vec[grid.padded_index(i, j, grid.kmt)] = 0.0f;
            }
        }
    }
}

// Vector operations
float dot_product(const float* x, const float* y, const GridParams& grid) {
    float result = 0.0f;
    
    // Compute dot product for inner grid points only
    for (int k = 0; k < grid.kmt; k++) {
        for (int j = 0; j < grid.jmt; j++) {
            for (int i = 0; i < grid.imt; i++) {
                int idx = grid.padded_index(i, j, k);
                result += x[idx] * y[idx];
            }
        }
    }
    
    return result;
}

float norm_squared(const float* x, const GridParams& grid) {
    return dot_product(x, x, grid);
}

} // namespace TensorMat19 